package com.learn.english.utils;

import com.learn.english.constants.APIConstant;
import com.learn.english.mapper.KnowledgeMapper;
import com.learn.english.model.entity.KnowledgeBase;
import com.learn.english.model.param.EmbeddingResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;

import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

/**
 * @author TISNIW
 * @version 1.0
 * @desc 知识检索工具类
 */
@Slf4j
@Component
public class RetrievalUtils {
    private final WebClient webClient;
    @Autowired
    private KnowledgeMapper knowledgeMapper;
    private static final float SIMILARITY_THRESHOLD = 0.5f;
    public RetrievalUtils() {
        this.webClient = WebClient.builder()
                .baseUrl(APIConstant.OLLAMA_MODEL_LIST_URL)
                .build();
    }
    /**
     * 生成文本的嵌入向量
     */
    public List<List<Double>> generateEmbeddings(List<String> inputs, String model) {
        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", model);
        requestBody.put("input", inputs);

        Mono<EmbeddingResponse> responseMono = webClient.post()
                .uri("/api/embed")
                .header("Content-Type", "application/json")
                .bodyValue(requestBody)
                .retrieve()
                .bodyToMono(EmbeddingResponse.class);

        EmbeddingResponse response = responseMono.block();
        return response != null ? response.getEmbeddings() : Collections.emptyList();
    }

    /**
     * 检索与问题最相似的知识条目
     */
    public List<KnowledgeBase> retrieveSimilarKnowledge(String query) {
        if (query == null || query.isEmpty()) {
            return Collections.emptyList();
        }

        List<KnowledgeBase> allEntries = knowledgeMapper.selectList();
        if (allEntries.isEmpty()) {
            return Collections.emptyList();
        }

        List<String> contents = allEntries.stream()
                .map(KnowledgeBase::getContent)
                .collect(Collectors.toList());

        // 将查询内容插入到第一个位置
        contents.add(0, query);

        List<List<Double>> embeddings = generateEmbeddings(contents, "bge-m3:567m");
        if (embeddings.isEmpty() || embeddings.size() != contents.size()) {
            return Collections.emptyList(); // 错误处理
        }

        List<Double> queryEmbedding = embeddings.get(0);
        List<List<Double>> entryEmbeddings = embeddings.subList(1, embeddings.size());

        return IntStream.range(0, allEntries.size())
                .boxed()
                .parallel()
                .filter(i -> computeCosineSimilarity(queryEmbedding, entryEmbeddings.get(i)) > SIMILARITY_THRESHOLD)
                .map(i -> {
                    KnowledgeBase knowledge = allEntries.get(i);
                    double similarity = computeCosineSimilarity(queryEmbedding, entryEmbeddings.get(i));
                    knowledge.setSimilarityScore(similarity);
                    return knowledge;
                })
                .sorted(Comparator.comparingDouble(KnowledgeBase::getSimilarityScore).reversed())
                .collect(Collectors.toList());
    }

    /**
     * 计算两个向量之间的余弦相似度
     */
    private double computeCosineSimilarity(List<Double> vecA, List<Double> vecB) {
        double dotProduct = 0.0, normA = 0.0, normB = 0.0;

        for (int i = 0; i < vecA.size(); i++) {
            double a = vecA.get(i);
            double b = vecB.get(i);
            dotProduct += a * b;
            normA += a * a;
            normB += b * b;
        }
        log.info("dotProduct: {}, normA: {}, normB: {}", dotProduct, normA, normB);
        log.info("similarity: {}", dotProduct / (Math.sqrt(normA) * Math.sqrt(normB)));
        return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
    }
}
